#include "Utils.h"

#define UtilsPoolTag 'TUHF' //File Hiding Utils Tag

NTSTATUS AllocateUnicodeString(OUT PUNICODE_STRING pUnicodeString,
                               IN const unsigned short ulLength,
                               IN const POOL_TYPE poolType)
{
    RtlZeroMemory(pUnicodeString, sizeof(UNICODE_STRING));

    pUnicodeString->Buffer = ExAllocatePoolWithTag(poolType, ulLength, UtilsPoolTag);

    if (NULL == pUnicodeString->Buffer)
    {
        return STATUS_INSUFFICIENT_RESOURCES;
    }

    pUnicodeString->MaximumLength = ulLength;

    return STATUS_SUCCESS;
}

void FreeUnicodeString(OUT PUNICODE_STRING pUnicodeString)
{
    if (NULL == pUnicodeString->Buffer)
    {
        return;
    }

    pUnicodeString->Length = 0;
    pUnicodeString->MaximumLength = 0;

    ExFreePoolWithTag(pUnicodeString->Buffer, UtilsPoolTag);

    pUnicodeString->Buffer = NULL;
}

NTSTATUS ConcatUnicodeString(IN const PUNICODE_STRING pusFirst,
                             IN const PUNICODE_STRING pusSecond,
                             OUT PUNICODE_STRING pusDestination,
                             IN const POOL_TYPE poolType)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;

    if (0 == pusFirst->Length)
    {
        return STATUS_INVALID_PARAMETER_1;
    }

    if (0 == pusSecond->Length)
    {
        return STATUS_INVALID_PARAMETER_2;
    }

    ntStatus = AllocateUnicodeString(pusDestination,
                                     pusFirst->Length + pusSecond->Length,
                                     poolType);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    ntStatus = RtlAppendUnicodeStringToString(pusDestination, pusFirst);

    ASSERT(NT_SUCCESS(ntStatus));

    ntStatus = RtlAppendUnicodeStringToString(pusDestination, pusSecond);

    ASSERT(NT_SUCCESS(ntStatus));

    return ntStatus;
}

// GetObjectName can be swapped to pagefile
#pragma alloc_text(PAGE, GetObjectName)

NTSTATUS GetObjectName(IN PVOID pObject,
                       OUT PUNICODE_STRING pusObjectName,
                       IN const POOL_TYPE poolType)
{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    unsigned long ulReturnLength = 0;

    PAGED_CODE();

    pusObjectName->Buffer = NULL;
    pusObjectName->Length = 0;
    pusObjectName->MaximumLength = 0;

    // get OBJECT_NAME_INFORMATION absolute size
    ntStatus = ObQueryNameString(pObject, NULL, 0, &ulReturnLength);

    if (STATUS_INFO_LENGTH_MISMATCH == ntStatus)
    {
        POBJECT_NAME_INFORMATION pObjectNameInfo = ExAllocatePoolWithTag(PagedPool, ulReturnLength, UtilsPoolTag);

        if (NULL == pObjectNameInfo)
        {
            return STATUS_INSUFFICIENT_RESOURCES;
        }

        ntStatus = ObQueryNameString(pObject, pObjectNameInfo, ulReturnLength, &ulReturnLength);

        if (!NT_SUCCESS(ntStatus))
        {
            ExFreePoolWithTag(pObjectNameInfo, UtilsPoolTag);

            return ntStatus;
        }

        // if ObQueryNameString() returned empty string, then go back
        if (0 == pObjectNameInfo->Name.Length)
        {
            return ntStatus;
        }
        
        // allocate pusObjectName and copy object name
        ntStatus = AllocateUnicodeString(pusObjectName, pObjectNameInfo->Name.Length, poolType);

        if (NT_SUCCESS(ntStatus))
        {
            RtlCopyUnicodeString(pusObjectName, &pObjectNameInfo->Name);
        }

        ExFreePoolWithTag(pObjectNameInfo, UtilsPoolTag);
    }

    return ntStatus;
}

void SetStatusIRP(IN OUT PIRP pIRP, IN const NTSTATUS ntStatus, IN const ULONG_PTR ulpInformation)
{
    pIRP->IoStatus.Status = ntStatus;
    pIRP->IoStatus.Information = ulpInformation;
}

NTSTATUS CompleteIRP(PIRP pIRP,
                     CCHAR ccPriorityBoost)
{
    NTSTATUS ntStatus = pIRP->IoStatus.Status;

    IoCompleteRequest(pIRP, ccPriorityBoost);

    return ntStatus;
}

NTSTATUS SetStatusAndCompleteIRP(PIRP pIRP,
                                 const NTSTATUS ntStatus,
                                 const ULONG_PTR ulpInformation,
                                 CCHAR ccPriorityBoost)
{
    SetStatusIRP(pIRP, ntStatus, ulpInformation);

    IoCompleteRequest(pIRP, ccPriorityBoost);

    return ntStatus;
}

NTSTATUS PassThroughIRP(PIRP pIRP, PDEVICE_OBJECT pNextLowerDevice)
{
    IoSkipCurrentIrpStackLocation(pIRP);

    return IoCallDriver(pNextLowerDevice, pIRP);
}

NTSTATUS WaitForIRPCompletion(IN PDEVICE_OBJECT pNextLowerDevice, IN PIRP pIRP)
{
    BOOLEAN bStatus = IoForwardIrpSynchronously(pNextLowerDevice, pIRP);

    if (!bStatus)
    {
        // no more device stacks
        return STATUS_UNSUCCESSFUL;
    }

    return STATUS_SUCCESS;
}

#pragma alloc_text(PAGE, EnumerateFSVolumes)

NTSTATUS EnumerateFSVolumes(IN PDEVICE_OBJECT pTargetCDO, IN PVOLUMEENUMERATOR pCallback, IN PVOID pContext)
{
    ULONG ulDevicesCount = 0;
    PDEVICE_OBJECT* ppDeviceList = NULL;
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    unsigned long ulDevicesCounter = 0;

    PAGED_CODE();

    // get volumes count
    ntStatus = IoEnumerateDeviceObjectList(pTargetCDO->DriverObject,
                                           NULL,
                                           0,
                                           &ulDevicesCount);

    if (NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    }

    // allocate memory for volume list
    ppDeviceList = ExAllocatePoolWithTag(PagedPool,
                                         ulDevicesCount * sizeof(PDEVICE_OBJECT),
                                         UtilsPoolTag);

    if (NULL == ppDeviceList)
    {

        return STATUS_INSUFFICIENT_RESOURCES;
    }

    // get volume list
    ntStatus = IoEnumerateDeviceObjectList(pTargetCDO->DriverObject,
                                           ppDeviceList,
                                           ulDevicesCount * sizeof(PDEVICE_OBJECT),
                                           &ulDevicesCount);

    if (!NT_SUCCESS(ntStatus))
    {
        ExFreePoolWithTag(ppDeviceList, UtilsPoolTag);

        return ntStatus;
    }

    // enumerate volumes
    for (ulDevicesCounter = 0; ulDevicesCounter < ulDevicesCount; ++ulDevicesCounter)
    {
        pCallback(ppDeviceList[ulDevicesCounter], pContext);

        ObDereferenceObject(ppDeviceList[ulDevicesCounter]);
    }

    ExFreePoolWithTag(ppDeviceList, UtilsPoolTag);

    return ntStatus;
}

#pragma alloc_text(PAGE, EnumerateMultiSZValue)

NTSTATUS EnumerateMultiSZValue(const PUNICODE_STRING pusKey,
                               const PUNICODE_STRING pusValue,
                               PENUMERATE_CALLBACK pCallback)

{
    NTSTATUS ntStatus = STATUS_UNSUCCESSFUL;
    HANDLE hKey = NULL;
    OBJECT_ATTRIBUTES objectAttributes = {0};
    PKEY_VALUE_PARTIAL_INFORMATION  pKeyValueInfo = NULL;
    unsigned long ulInformationSize = 0;
    wchar_t* wcBufStart = NULL;
    wchar_t* wcBufEnd = NULL;

    PAGED_CODE();


    InitializeObjectAttributes(&objectAttributes, pusKey, OBJ_KERNEL_HANDLE, NULL, NULL);

    ntStatus = ZwOpenKey(&hKey, KEY_ALL_ACCESS, &objectAttributes);

    if (!NT_SUCCESS(ntStatus))
    {
        return ntStatus;
    } 

    ntStatus = ZwQueryValueKey(hKey, pusValue, KeyValuePartialInformation, NULL, 0, &ulInformationSize);

    if (0 == ulInformationSize)
    {
        ZwClose(hKey);

        return ntStatus;
    }  


    pKeyValueInfo =(KEY_VALUE_PARTIAL_INFORMATION*)ExAllocatePoolWithTag(PagedPool, ulInformationSize, UtilsPoolTag);

    if (NULL == pKeyValueInfo)  
    {
        ZwClose(hKey);

        return STATUS_INSUFFICIENT_RESOURCES;
    }

    ntStatus = ZwQueryValueKey(hKey, pusValue, KeyValuePartialInformation, pKeyValueInfo, ulInformationSize, &ulInformationSize);

    if (!NT_SUCCESS(ntStatus))
    {
        ExFreePoolWithTag(pKeyValueInfo, UtilsPoolTag);
        ZwClose(hKey);

        return ntStatus;
    }

    if (REG_MULTI_SZ != pKeyValueInfo->Type)
    {
        ExFreePoolWithTag(pKeyValueInfo, UtilsPoolTag);
        ZwClose(hKey);

        return STATUS_INVALID_PARAMETER_2;
    } 

    wcBufStart = (wchar_t*)&(pKeyValueInfo->Data[0]);
    wcBufEnd = wcBufStart + pKeyValueInfo->DataLength - 2 * sizeof(wchar_t) /*\0\0 at the end of MULTI_SZ value*/;

    while(wcBufStart < wcBufEnd)
    {
        ntStatus = pCallback(pusValue, wcBufStart);

        if (!NT_SUCCESS(ntStatus))
        {
            ExFreePoolWithTag(pKeyValueInfo, UtilsPoolTag);
            ZwClose(hKey);

            return ntStatus;
        }

        wcBufStart += wcslen(wcBufStart) + 1;
    }

    ExFreePoolWithTag(pKeyValueInfo, UtilsPoolTag);

    ZwClose(hKey);

    return STATUS_SUCCESS;
}

NTSTATUS AcquireResource(IN PERESOURCE pResource, IN const BOOLEAN bExclusive)
{
    BOOLEAN bStatus = FALSE;

    KeEnterCriticalRegion();

    if (bExclusive)
    {
        bStatus = ExAcquireResourceExclusiveLite(pResource, TRUE);
    }
    else
    {
        bStatus = ExAcquireResourceSharedLite(pResource, TRUE);
    }

    if (!bStatus)
    {
        KeLeaveCriticalRegion();

        return STATUS_UNSUCCESSFUL;
    }

    return STATUS_SUCCESS;
}

NTSTATUS ReleaseResource(IN PERESOURCE pResource)
{
    if (!ExIsResourceAcquiredExclusiveLite(pResource) && !ExIsResourceAcquiredSharedLite(pResource))
    {
        return STATUS_RESOURCE_NOT_OWNED;
    }

    ExReleaseResourceLite(pResource);

    KeLeaveCriticalRegion();

    return STATUS_SUCCESS;
}
